import matplotlib.pyplot as plt
from models import *
import torchvision.transforms as transforms
from  PIL import Image
from config import *

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = CNN(class_num).to(device)

model.load_state_dict(torch.load("./network.pth"))

tranform_img=transforms.Compose([transforms.Resize((224,224)),
        transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])

image_name="./testset.jpg"
image = Image.open(image_name).convert('RGB')
x = tranform_img(image)
x =x.unsqueeze(0).to(device)  #必须加批次1

with torch.no_grad():
        output_class = model(x)

print(output_class)

plt.figure(figsize=(16,8))
plt.imshow(image)
plt.text(0, -10, output_class, fontsize=20 , color='red', wrap=True)
plt.show()



